package io.gitee.luoyuan800.tool.llm.web;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.Spliterator;
import java.util.Spliterators;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CompletionException;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Flow;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

/**
 * 基于java 的实现
 * @author Luoyuan/ Gavin
 */
@Slf4j
public class JavaHttpClient implements RestfulClient {
    private final HttpClient httpClient;
    private final ObjectMapper objectMapper;

    public JavaHttpClient() {
        this(HttpClient.newBuilder()
                        .connectTimeout(Duration.ofSeconds(10))
                        .build(),
                new ObjectMapper());
    }

    public JavaHttpClient(ObjectMapper objectMapper) {
        this(HttpClient.newBuilder()
                        .connectTimeout(Duration.ofSeconds(10))
                        .build(),
                objectMapper);
    }

    public JavaHttpClient(HttpClient httpClient, ObjectMapper objectMapper) {
        this.httpClient = httpClient;
        this.objectMapper = objectMapper;
    }

    private static HttpRequest buildRequest(String url, Map<String, String> header, String jsonBody, boolean stream) {
        HttpRequest.Builder builder = HttpRequest.newBuilder()
                .uri(URI.create(url))
                .header("Content-Type", "application/json");
        if (header != null) {
            header.forEach(builder::header);
        }
        if (stream) {
            builder.header("Accept", "text/event-stream");
        }
        return builder
                .POST(HttpRequest.BodyPublishers.ofString(jsonBody, StandardCharsets.UTF_8))
                .build();
    }

    private <R> R readValueAsObject(String json, Class<R> responseType) {
        try {
            if(StringUtils.equalsAnyIgnoreCase(json, "data:[DONE]", "data: [DONE]")){
                return null;
            }
            return objectMapper.readValue(json.replaceFirst("^data:", ""), responseType);
        } catch (JsonProcessingException e) {
            log.error("Failed to process json {}.", json, e);
            return null;
        }
    }

    public <T, R> Stream<R> postStream(String url, Map<String, String> header, T requestBody, Class<R> responseType) {
        try {
            // 创建阻塞队列作为数据缓冲区
            SuppressExceptionLinkedBlockingQueue<R> queue = new SuppressExceptionLinkedBlockingQueue<>();
            AtomicReference<Throwable> error = new AtomicReference<>();
            CountDownLatch completed = new CountDownLatch(1);
            String jsonBody = objectMapper.writeValueAsString(requestBody);
            HttpRequest request = buildRequest(url, header, jsonBody, true);
            httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofPublisher()).thenApply(HttpResponse::body)
                    .thenAccept(publisher -> publisher.subscribe(new Flow.Subscriber<>() {
                        private final StringBuilder buffer = new StringBuilder();
                        private Flow.Subscription subscription;

                        @Override
                        public void onSubscribe(Flow.Subscription subscription) {
                            this.subscription = subscription;
                            subscription.request(1); // 初始请求一个数据块
                        }

                        @Override
                        public void onNext(List<ByteBuffer> item) {
                            try {
                                // 将字节缓冲区转换为字符串
                                for (ByteBuffer byteBuffer : item) {
                                    String chunk = StandardCharsets.UTF_8.decode(byteBuffer).toString();
                                    buffer.append(chunk);

                                    // 处理可能的多行数据
                                    int newlinePos;
                                    while ((newlinePos = buffer.indexOf("\n")) != -1) {
                                        String line = buffer.substring(0, newlinePos).trim();
                                        buffer.delete(0, newlinePos + 1);

                                        if (!line.isEmpty()) {
                                            queue.add(readValueAsObject(line, responseType));
                                        }
                                    }
                                }
                                subscription.request(1); // 请求下一个数据块
                            } catch (Exception e) {
                                log.error("Failed to read value from stream response.", e);
                            }
                        }

                        @Override
                        public void onError(Throwable throwable) {
                            error.set(throwable);
                            completed.countDown();
                        }

                        @Override
                        public void onComplete() {
                            // 处理缓冲区中剩余的数据
                            if (!buffer.isEmpty()) {
                                queue.add(readValueAsObject(buffer.toString(), responseType));
                            }
                            completed.countDown();
                        }


                    }))
                    .exceptionally(e -> {
                        error.set(e);
                        completed.countDown();
                        return null;
                    });
            return creatStream(queue, error, completed);
        } catch (JsonProcessingException e) {
            throw new RuntimeException(e);
        }
    }

    private <R> Stream<R> creatStream(BlockingQueue<R> queue, AtomicReference<Throwable> error, CountDownLatch completed) {
        // 创建并返回Stream
        // Stream关闭时的处理
        return StreamSupport.stream(new Spliterators.AbstractSpliterator<R>(
                Long.MAX_VALUE, Spliterator.ORDERED | Spliterator.NONNULL) {

            @Override
            public boolean tryAdvance(Consumer<? super R> action) {
                try {
                    while (true) {
                        // 检查是否完成或有错误
                        if (completed.getCount() == 0) {
                            if (error.get() != null) {
                                throw new CompletionException(error.get());
                            }
                            R lastItem = queue.poll();
                            if (lastItem != null) {
                                action.accept(lastItem);
                            }
                            return false;
                        }

                        // 从队列获取数据
                        R item = queue.poll(100, TimeUnit.MILLISECONDS);
                        if (item != null) {
                            action.accept(item);
                            return true;
                        }
                    }
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    throw new CompletionException(e);
                }
            }
        }, false).onClose(completed::countDown);
    }

    /**
     * POST请求，将请求体序列化为JSON，并将响应反序列化为指定类型
     *
     * @param url          请求URL
     * @param requestBody  请求体对象
     * @param responseType 响应类型
     * @param <T>          请求体类型
     * @param <R>          响应体类型
     * @return 响应对象
     */
    public <T, R> R post(String url, Map<String, String> header, T requestBody, Class<R> responseType) {
        try {
            String jsonBody = objectMapper.writeValueAsString(requestBody);
            HttpRequest request = buildRequest(url, header, jsonBody, false);
            HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());

            if (response.statusCode() >= 200 && response.statusCode() < 300) {
                return objectMapper.readValue(response.body(), responseType);
            } else {
                throw new RestfulClientException("HTTP request failed with status: " + response.statusCode() + ", body: " + response.body());
            }
        } catch (JsonProcessingException e) {
            throw new RestfulClientException("Failed to serialize request body", e);
        } catch (IOException | InterruptedException e) {
            throw new RestfulClientException("HTTP request failed", e);
        }
    }

    /**
     * POST请求，支持带URI变量
     *
     * @param url          请求URL模板
     * @param requestBody  请求体对象
     * @param responseType 响应类型
     * @param uriVariables URI变量
     * @param <T>          请求体类型
     * @param <R>          响应体类型
     * @return 响应对象
     */
    public <T, R> R post(String url, Map<String, String> header, T requestBody, Class<R> responseType, Map<String, Object> uriVariables) {
        String expandedUrl = expandUrl(url, uriVariables);
        return post(expandedUrl, header, requestBody, responseType);
    }

    private String expandUrl(String url, Map<String, Object> uriVariables) {
        if (uriVariables == null || uriVariables.isEmpty()) {
            return url;
        }

        String result = url;
        for (Map.Entry<String, Object> entry : uriVariables.entrySet()) {
            String placeholder = "{" + entry.getKey() + "}";
            String value = entry.getValue().toString();
            result = result.replace(placeholder, value);
        }
        return result;
    }

}
